# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import functools
import ujson as json

from coin.base.book.types import TradeSide, create_market_trade_event
from coin.exchange.base.ws.handler import WebsocketHandler
from coin.exchange.kraken.book_builder import RestBookBuilder as BookBuilder
from coin.exchange.kraken.kr_rest.product import (KrakenProduct, remove_xz_from_rest_native_symbol)


@functools.lru_cache(maxsize=32)
def get_product_from_symbol(symbol):
  return KrakenProduct.FromStrNativeProduct(symbol)


class RestHandlerAsOfWebsocket(WebsocketHandler):
  pass


class KrakenRestHandler(RestHandlerAsOfWebsocket):
  def __init__(self, *, feed_checker_config=None):
    super().__init__(feed_checker_config=feed_checker_config)
    self._last_trades = {}

  def on_depth(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())
    error = data['error']
    if error:
      return

    symbols = list(data['result'])
    assert len(symbols) == 1
    symbol = symbols[0]
    data = data['result'][symbol]
    symbol = remove_xz_from_rest_native_symbol(symbol)
    if not self.is_subscribed('book', symbol):
      return

    if symbol not in self._books:
      book = BookBuilder(symbol, book_builder_observer=self._book_builder_observer)
      self._books[symbol] = book

    else:
      book = self._books[symbol]
      book.clear_book()

    assert 'bids' in data
    assert 'asks' in data
    bids = data['bids']
    asks = data['asks']
    book.snapshot(bids, asks)
    book.update_timestamp(record.timestamp)
    book.update_topic_id(record.topic_id)
    self.ready = True
    product = get_product_from_symbol(symbol)
    if self._feed_checker:
      self._feed_checker.on_book(product=product, book=book, timestamp=record.timestamp)
    self.publish_event('book', book.symbol, book)

  def on_trade(self, record, queue_data, topic_data):
    data = record.data
    if not isinstance(data, dict):
      data = json.loads(data.tobytes())
    error = data['error']
    if error:
      return

    if 'last' in data['result']:
      last_timestamp = int(data['result']['last'])
    else:
      return
    data = data['result']
    data.pop('last')
    symbols = list(data)
    assert len(symbols) == 1
    symbol = symbols[0]
    data = data[symbol]
    symbol = remove_xz_from_rest_native_symbol(symbol)
    if not self.is_subscribed('trade', symbol):
      return

    self._last_trades.setdefault(symbol, 0)
    last_trade_sec = self._last_trades[symbol] / 1e9

    assert isinstance(data, list)
    for event in data:
      price, qty, time_sec, side, _, _ = event
      if time_sec <= last_trade_sec:
        continue
      price = float(price)
      qty = float(qty)
      assert side in ['b', 's']
      side = TradeSide.TRADE_SELL_SIDE if side == 's' else TradeSide.TRADE_BUY_SIDE
      trade_event = create_market_trade_event(record.timestamp,
                                              symbol,
                                              price,
                                              qty,
                                              side,
                                              topic_id=record.topic_id)
      product = get_product_from_symbol(symbol)
      if self._feed_checker:
        self._feed_checker.on_trade(product=product,
                                    trade_event=trade_event,
                                    timestamp=record.timestamp)
      self.publish_event('trade', symbol, trade_event)

    self._last_trades[symbol] = last_timestamp
